package ai

import (
	"encoding/json"
	"fmt"
	"net/http"
	"project-info/src/handle/mcp"
	"project-info/src/model"
	"sync"
	"time"
	"github.com/gorilla/websocket"

	"project-info/src/lib")

// AI WebSocket消息类型
type AIWSMessage struct {
	Type      string                 `json:"type"`      // 消息类型: user_message, ai_response, tool_call, tool_result, error, session_info
	ID        string                 `json:"id"`        // 消息ID
	SessionID string                 `json:"sessionId"` // 会话ID
	Content   string                 `json:"content"`   // 消息内容
	Data      map[string]interface{} `json:"data"`      // 附加数据
	Timestamp int64                  `json:"timestamp"` // 时间戳
}

// AI会话状态
type AISession struct {
	ID              string                   `json:"id"`
	UserID          string                   `json:"userId"`
	Messages        []Message                `json:"messages"`        // 对话历史
	Context         map[string]interface{}   `json:"context"`         // 会话上下文
	PendingActions  []PendingAction          `json:"pendingActions"`  // 待执行的操作
	AvailableTools  []mcp.MCPTool            `json:"availableTools"`  // 可用工具
	ToolsByCategory map[string][]mcp.MCPTool `json:"toolsByCategory"` // 按分类的工具
	CurrentTask     *TaskInfo                `json:"currentTask"`     // 当前任务信息
	TaskHistory     []TaskInfo               `json:"taskHistory"`     // 任务历史
	CreatedAt       time.Time                `json:"createdAt"`
	UpdatedAt       time.Time                `json:"updatedAt"`
	mutex           sync.RWMutex             `json:"-"`
}

// 任务信息
type TaskInfo struct {
	ID        string                 `json:"id"`
	UserInput string                 `json:"userInput"` // 用户输入
	Status    string                 `json:"status"`    // pending, processing, completed, failed
	Messages  []Message              `json:"messages"`  // 任务相关的消息
	ToolCalls []ToolCallInfo         `json:"toolCalls"` // 工具调用记录
	Result    string                 `json:"result"`    // 最终结果
	StartTime time.Time              `json:"startTime"`
	EndTime   *time.Time             `json:"endTime"`
	Context   map[string]interface{} `json:"context"` // 任务上下文
}

// 工具调用信息
type ToolCallInfo struct {
	ToolName   string                 `json:"toolName"`
	Parameters map[string]interface{} `json:"parameters"`
	Result     *mcp.MCPToolResult     `json:"result"`
	Timestamp  time.Time              `json:"timestamp"`
}

// 待执行操作
type PendingAction struct {
	ID         string                 `json:"id"`
	ToolName   string                 `json:"toolName"`
	Parameters map[string]interface{} `json:"parameters"`
	Status     string                 `json:"status"` // pending, executing, completed, failed
	Result     *mcp.MCPToolResult     `json:"result"`
	CreatedAt  time.Time              `json:"createdAt"`
}

// AI WebSocket连接管理
type AIWSConnection struct {
	ws        *websocket.Conn
	session   *AISession
	mcpClient *mcp.MCPInternalClient
	context   *lib.GinContext
	closed    bool
	mutex     sync.RWMutex
}

// 全局会话管理器
var (
	aiSessions     = make(map[string]*AISession)
	aiSessionMutex = sync.RWMutex{}
)

// WebSocket升级器
var aiUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	Subprotocols: []string{"ai-chat"},
}

// AI WebSocket处理器
func AIWebSocketHandler(c *lib.GinContext) error {
	ws, err := aiUpgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		c.Logger().Error("AI WebSocket升级失败: " + err.Error())
		return err
	}
	defer ws.Close()

	c.Logger().Info("新的AI WebSocket客户端连接")

	// 创建MCP内部客户端
	mcpClient := mcp.NewMCPInternalClient(c)

	// 创建WebSocket连接
	conn := &AIWSConnection{
		ws:        ws,
		mcpClient: mcpClient,
		context:   c,
		closed:    false,
	}

	// 发送初始化信息
	if err := conn.sendInitialInfo(); err != nil {
		c.Logger().Error("发送初始化信息失败: " + err.Error())
		return err
	}

	// 处理消息循环
	for {
		if conn.isClosed() {
			break
		}

		_, messageBytes, err := ws.ReadMessage()
		if err != nil {
			c.Logger().Error("读取AI WebSocket消息失败: " + err.Error())
			break
		}

		var msg AIWSMessage
		if err := json.Unmarshal(messageBytes, &msg); err != nil {
			c.Logger().Error("解析AI WebSocket消息失败: " + err.Error())
			continue
		}

		// 处理消息
		go conn.handleMessage(&msg)
	}

	// 清理会话
	if conn.session != nil {
		conn.cleanupSession()
	}

	return nil
}

// 发送初始化信息
func (conn *AIWSConnection) sendInitialInfo() error {
	availableTools := conn.mcpClient.GetAvailableTools()
	toolsByCategory := conn.mcpClient.GetToolsByCategory()

	initMsg := AIWSMessage{
		Type:      "session_info",
		ID:        generateMessageID(),
		Content:   "AI助手已准备就绪",
		Timestamp: time.Now().Unix(),
		Data: map[string]interface{}{
			"availableTools":  availableTools,
			"toolsByCategory": toolsByCategory,
			"capabilities": map[string]interface{}{
				"multiTurn":     true,
				"toolChaining":  true,
				"contextMemory": true,
			},
		},
	}

	return conn.sendMessage(&initMsg)
}

// 处理WebSocket消息
func (conn *AIWSConnection) handleMessage(msg *AIWSMessage) {
	switch msg.Type {
	case "user_message":
		conn.handleUserMessage(msg)
	case "create_session":
		conn.handleCreateSession(msg)
	case "get_session":
		conn.handleGetSession(msg)
	case "clear_session":
		conn.handleClearSession(msg)
	default:
		conn.sendError(msg.ID, "未知消息类型: "+msg.Type)
	}
}

// 处理用户消息
func (conn *AIWSConnection) handleUserMessage(msg *AIWSMessage) {
	if conn.session == nil {
		conn.sendError(msg.ID, "请先创建会话")
		return
	}

	// 检查是否有正在进行的任务
	if conn.session.hasActiveTask() {
		fmt.Printf("handleUserMessage, hasActiveTask: true\n")
		conn.sendError(msg.ID, "上一个任务正在处理中，请等待完成后再发送新消息")
		return
	}

	// 创建新任务
	taskID := generateTaskID()
	task := &TaskInfo{
		ID:        taskID,
		UserInput: msg.Content,
		Status:    "processing",
		Messages:  []Message{{Role: "user", Content: msg.Content}},
		ToolCalls: []ToolCallInfo{},
		StartTime: time.Now(),
		Context:   make(map[string]interface{}),
	}

	// 设置当前任务
	conn.session.setCurrentTask(task)

	// 开始任务处理
	go conn.processTask(msg, task)
}

// 处理任务（基于任务的多轮对话）
func (conn *AIWSConnection) processTask(initialMsg *AIWSMessage, task *TaskInfo) {
	defer func() {
		// 任务完成后的清理工作
		conn.session.completeCurrentTask()

		// 发送任务状态更新，确保前端知道任务已结束
		taskStatusMsg := AIWSMessage{
			Type:      "task_status_update",
			ID:        generateMessageID(),
			SessionID: conn.session.ID,
			Content:   "任务处理结束",
			Data: map[string]interface{}{
				"taskId":            task.ID,
				"status":            task.Status,
				"canSendNewMessage": true,
			},
			Timestamp: time.Now().Unix(),
		}
		conn.sendMessage(&taskStatusMsg)
	}()

	maxTurns := model.C.AI.MaxTurns // 从配置文件读取最大对话轮数
	if maxTurns <= 0 {
		maxTurns = 10 // 默认值
	}
	currentTurn := 0

	for currentTurn < maxTurns {
		fmt.Printf("currentTurn: %d\n", currentTurn)
		currentTurn++

		// 构建基于当前任务的系统提示词
		systemPrompt := conn.buildTaskBasedSystemPrompt(task)

		// 调用AI模型（只使用当前任务的消息历史）
		aiResponse, err := conn.callAIModelWithTaskHistory(systemPrompt, task)
		if err != nil {
			task.Status = "failed"
			conn.sendError(initialMsg.ID, "调用AI模型失败: "+err.Error())
			return
		}

		// 添加AI响应到任务消息历史
		task.addMessage("assistant", aiResponse)

		// 发送AI响应给客户端
		responseMsg := AIWSMessage{
			Type:      "ai_response",
			ID:        generateMessageID(),
			SessionID: conn.session.ID,
			Content:   aiResponse,
			Timestamp: time.Now().Unix(),
			Data: map[string]interface{}{
				"turn":   currentTurn,
				"taskId": task.ID,
			},
		}
		conn.sendMessage(&responseMsg)

		// 解析是否需要调用工具
		toolName, params := parseAIResponseForMCP(aiResponse)
		if toolName == "" {
			// 没有工具调用，任务完成
			task.Status = "completed"
			task.Result = aiResponse
			now := time.Now()
			task.EndTime = &now

			// 通知前端任务完成
			taskCompleteMsg := AIWSMessage{
				Type:      "task_completed",
				ID:        generateMessageID(),
				SessionID: conn.session.ID,
				Content:   "任务处理完成",
				Data: map[string]interface{}{
					"taskId":  task.ID,
					"status":  "completed",
					"endTime": task.EndTime,
				},
				Timestamp: time.Now().Unix(),
			}
			conn.sendMessage(&taskCompleteMsg)
			break
		}

		// 执行工具调用
		if !conn.executeTaskToolCall(task, toolName, params, currentTurn) {
			// 工具调用失败，任务失败
			task.Status = "failed"
			now := time.Now()
			task.EndTime = &now

			// 通知前端任务失败
			taskFailedMsg := AIWSMessage{
				Type:      "task_failed",
				ID:        generateMessageID(),
				SessionID: conn.session.ID,
				Content:   "任务处理失败",
				Data: map[string]interface{}{
					"taskId":  task.ID,
					"status":  "failed",
					"endTime": task.EndTime,
				},
				Timestamp: time.Now().Unix(),
			}
			conn.sendMessage(&taskFailedMsg)
			break
		}
	}

	if currentTurn >= maxTurns {
		task.Status = "failed"
		now := time.Now()
		task.EndTime = &now

		// 通知前端任务因轮数限制失败
		taskFailedMsg := AIWSMessage{
			Type:      "task_failed",
			ID:        generateMessageID(),
			SessionID: conn.session.ID,
			Content:   "任务处理轮数超过限制，任务终止",
			Data: map[string]interface{}{
				"taskId":  task.ID,
				"status":  "failed",
				"reason":  "max_turns_exceeded",
				"endTime": task.EndTime,
			},
			Timestamp: time.Now().Unix(),
		}
		conn.sendMessage(&taskFailedMsg)
	}
}

// 执行任务工具调用
func (conn *AIWSConnection) executeTaskToolCall(task *TaskInfo, toolName string, params map[string]interface{}, turn int) bool {
	// 发送工具调用通知
	// toolCallMsg := AIWSMessage{
	// 	Type:      "tool_call",
	// 	ID:        generateMessageID(),
	// 	SessionID: conn.session.ID,
	// 	Content:   fmt.Sprintf("正在执行工具: %s", toolName),
	// 	Timestamp: time.Now().Unix(),
	// 	Data: map[string]interface{}{
	// 		"toolName":   toolName,
	// 		"parameters": params,
	// 		"turn":       turn,
	// 		"taskId":     task.ID,
	// 	},
	// }
	// conn.sendMessage(&toolCallMsg)

	// 执行工具
	toolResult := conn.mcpClient.CallTool(toolName, params)

	// 记录工具调用信息
	toolCallInfo := ToolCallInfo{
		ToolName:   toolName,
		Parameters: params,
		Result:     toolResult,
		Timestamp:  time.Now(),
	}
	task.addToolCall(toolCallInfo)

	// 发送工具执行结果
	// toolResultMsg := AIWSMessage{
	// 	Type:      "tool_result",
	// 	ID:        generateMessageID(),
	// 	SessionID: conn.session.ID,
	// 	Content:   toolResult.Content,
	// 	Timestamp: time.Now().Unix(),
	// 	Data: map[string]interface{}{
	// 		"toolName": toolName,
	// 		"success":  toolResult.Success,
	// 		"result":   toolResult,
	// 		"turn":     turn,
	// 		"taskId":   task.ID,
	// 	},
	// }
	// conn.sendMessage(&toolResultMsg)

	if !toolResult.Success {
		// 工具执行失败，添加错误信息到任务上下文
		errorContext := fmt.Sprintf("工具 %s 执行失败: %s", toolName, toolResult.Error)
		task.addMessage("system", errorContext)
		return false
	}

	// 工具执行成功，添加结果到任务上下文
	successContext := fmt.Sprintf("工具 %s 执行成功，结果: %s", toolName, toolResult.Content)
	task.addMessage("system", successContext)

	// 更新任务上下文
	task.updateContext(toolName, toolResult)

	return true
}

// 构建基于任务的系统提示词
func (conn *AIWSConnection) buildTaskBasedSystemPrompt(task *TaskInfo) string {
	// availableTools := conn.mcpClient.GetAvailableTools()
	// toolDescriptions := make([]string, 0, len(availableTools))
	// for _, tool := range availableTools {
	// 	inputSchema, _ := json.Marshal(tool.InputSchema)
	// 	toolDescriptions = append(toolDescriptions, fmt.Sprintf("- %s: %s %s", tool.Name, tool.Description, inputSchema))
	// }

	// fmt.Println(toolDescriptions)

	// 获取当前任务的上下文信息
	// contextInfo := ""
	// if len(task.Context) > 0 {
	// 	contextInfo = "\n\n当前任务上下文:\n"
	// 	for key, value := range task.Context {
	// 		contextInfo += fmt.Sprintf("- %s: %v\n", key, value)
	// 	}
	// }

	// // 获取已执行的工具调用信息
	// toolCallsInfo := ""
	// if len(task.ToolCalls) > 0 {
	// 	toolCallsInfo = "\n\n本任务已执行的工具调用：\n"
	// 	for i, toolCall := range task.ToolCalls {
	// 		status := "成功"
	// 		if toolCall.Result != nil && !toolCall.Result.Success {
	// 			status = "失败"
	// 		}
	// 		toolCallsInfo += fmt.Sprintf("%d. %s (%s)\n", i+1, toolCall.ToolName, status)
	// 	}
	// }
	// 	你可以使用以下工具来完成这个任务：
	// %s

	return fmt.Sprintf(`你是一个智能运维助手，正在处理用户的单个任务。你需要专注于完成当前任务，不要参考其他任务的历史信息。

用户任务：%s

任务处理规则：
1. 专注于当前任务，分析用户的具体需求
2. 根据需要调用相关工具获取信息
3. 基于工具执行结果决定下一步操作
4. 工具之间可能存在依赖关系，请按逻辑顺序执行
5. 任务完成后提供清晰的总结
6. 如判断无需调用工具,那么返回的结果整理成方便html页面展示的格式

工具调用格式：
TOOL_CALL: tool_name
TOOL_PARAMS: {"param1": "value1", "param2": "value2"}

注意事项：
- 只调用完成当前任务必需的工具
- 参数必须是有效的JSON格式
- 优先使用任务上下文中的信息
- 工具名称,描述后跟的是可用参数信息,properties中定义的是参数结构,required表示是否必填,description表示参数描述,type表示参数类型
- 工具调用结果中包含的信息，直接用于任务上下文，无需再次调用工具
- 如果工具执行失败，分析原因并尝试替代方案

请用中文回复，保持专业和友好的语调。`, task.UserInput)
}

// 使用任务历史调用AI模型
func (conn *AIWSConnection) callAIModelWithTaskHistory(systemPrompt string, task *TaskInfo) (string, error) {
	// 构建消息历史，包含系统提示词
	messages := []Message{{Role: "system", Content: systemPrompt}}

	// 添加当前任务的消息历史（只使用当前任务相关的消息）
	maxHistoryMessages := model.C.AI.MaxHistoryMessages
	if maxHistoryMessages <= 0 {
		maxHistoryMessages = 20 // 默认值
	}
	taskMessages := task.getRecentMessages(maxHistoryMessages)
	messages = append(messages, taskMessages...)

	availableTools := conn.mcpClient.GetAvailableTools()

	reqBody := OpenAIRequest{
		Model:       model.C.AI.Model,
		Messages:    messages,
		MaxTokens:   model.C.AI.MaxTokens,
		Temperature: model.C.AI.Temperature,
		Tools:       availableTools,
	}

	return callAIModelWithRequest(reqBody)
}

// 会话管理方法
func (conn *AIWSConnection) handleCreateSession(msg *AIWSMessage) {
	userID := "default" // 这里应该从认证信息中获取用户ID
	if uid, ok := msg.Data["userId"].(string); ok {
		userID = uid
	}

	session := &AISession{
		ID:              generateSessionID(),
		UserID:          userID,
		Messages:        []Message{},
		Context:         make(map[string]interface{}),
		PendingActions:  []PendingAction{},
		AvailableTools:  conn.mcpClient.GetAvailableTools(),
		ToolsByCategory: conn.mcpClient.GetToolsByCategory(),
		CurrentTask:     nil,
		TaskHistory:     []TaskInfo{},
		CreatedAt:       time.Now(),
		UpdatedAt:       time.Now(),
	}

	// 保存会话
	aiSessionMutex.Lock()
	aiSessions[session.ID] = session
	aiSessionMutex.Unlock()

	conn.session = session

	// 发送会话创建成功消息
	responseMsg := AIWSMessage{
		Type:      "session_created",
		ID:        generateMessageID(),
		SessionID: session.ID,
		Content:   "会话创建成功，支持任务隔离的多轮对话",
		Timestamp: time.Now().Unix(),
		Data: map[string]interface{}{
			"session": session,
			"features": map[string]bool{
				"taskIsolation":  true,
				"contextMemory":  true,
				"toolChaining":   true,
				"concurrentTask": false,
			},
		},
	}
	conn.sendMessage(&responseMsg)
}

func (conn *AIWSConnection) handleGetSession(msg *AIWSMessage) {
	sessionID, ok := msg.Data["sessionId"].(string)
	if !ok {
		conn.sendError(msg.ID, "缺少sessionId参数")
		return
	}

	aiSessionMutex.RLock()
	session, exists := aiSessions[sessionID]
	aiSessionMutex.RUnlock()

	if !exists {
		conn.sendError(msg.ID, "会话不存在")
		return
	}

	conn.session = session

	responseMsg := AIWSMessage{
		Type:      "session_loaded",
		ID:        generateMessageID(),
		SessionID: session.ID,
		Content:   "会话加载成功",
		Timestamp: time.Now().Unix(),
		Data: map[string]interface{}{
			"session": session,
		},
	}
	conn.sendMessage(&responseMsg)
}

func (conn *AIWSConnection) handleClearSession(msg *AIWSMessage) {
	if conn.session == nil {
		conn.sendError(msg.ID, "没有活动会话")
		return
	}

	// 检查是否有正在进行的任务
	if conn.session.hasActiveTask() {
		conn.sendError(msg.ID, "有任务正在进行中，无法清空会话")
		return
	}

	// 清空会话历史和上下文
	conn.session.mutex.Lock()
	conn.session.Messages = []Message{}
	conn.session.Context = make(map[string]interface{})
	conn.session.PendingActions = []PendingAction{}
	conn.session.CurrentTask = nil
	conn.session.TaskHistory = []TaskInfo{}
	conn.session.UpdatedAt = time.Now()
	conn.session.mutex.Unlock()

	responseMsg := AIWSMessage{
		Type:      "session_cleared",
		ID:        generateMessageID(),
		SessionID: conn.session.ID,
		Content:   "会话已清空，包括所有任务历史",
		Timestamp: time.Now().Unix(),
	}
	conn.sendMessage(&responseMsg)
}

// 辅助方法
func (conn *AIWSConnection) sendMessage(msg *AIWSMessage) error {
	conn.mutex.RLock()
	defer conn.mutex.RUnlock()

	if conn.closed {
		return fmt.Errorf("连接已关闭")
	}

	return conn.ws.WriteJSON(msg)
}

func (conn *AIWSConnection) sendError(msgID, errorMsg string) {
	errorResponse := AIWSMessage{
		Type:      "error",
		ID:        msgID,
		Content:   errorMsg,
		Timestamp: time.Now().Unix(),
	}
	conn.sendMessage(&errorResponse)
}

func (conn *AIWSConnection) isClosed() bool {
	conn.mutex.RLock()
	defer conn.mutex.RUnlock()
	return conn.closed
}

func (conn *AIWSConnection) cleanupSession() {
	conn.mutex.Lock()
	conn.closed = true
	conn.mutex.Unlock()

	if conn.session != nil {
		// 这里可以实现会话持久化逻辑
		conn.context.Logger().Info("清理AI会话: " + conn.session.ID)
	}
}

// AISession 方法
func (s *AISession) hasActiveTask() bool {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	fmt.Printf("hasActiveTask, CurrentTask: %v\n", s.CurrentTask)
	return s.CurrentTask != nil && (s.CurrentTask.Status == "processing" || s.CurrentTask.Status == "pending")
}

func (s *AISession) setCurrentTask(task *TaskInfo) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.CurrentTask = task
	s.UpdatedAt = time.Now()
}

func (s *AISession) completeCurrentTask() {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.CurrentTask != nil {
		// 将当前任务添加到历史记录
		s.TaskHistory = append(s.TaskHistory, *s.CurrentTask)

		// 限制历史记录数量
		maxHistory := 50 // 最多保留50个任务历史
		if len(s.TaskHistory) > maxHistory {
			s.TaskHistory = s.TaskHistory[len(s.TaskHistory)-maxHistory:]
		}

		// 清空当前任务
		s.CurrentTask = nil
	}
	s.UpdatedAt = time.Now()
}

func (s *AISession) addMessage(role, content string) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.Messages = append(s.Messages, Message{
		Role:    role,
		Content: content,
	})
	s.UpdatedAt = time.Now()
}

func (s *AISession) getRecentMessages(limit int) []Message {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	if len(s.Messages) <= limit {
		return s.Messages
	}
	return s.Messages[len(s.Messages)-limit:]
}

func (s *AISession) updateContext(key string, value interface{}) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.Context[key] = value
	s.UpdatedAt = time.Now()
}

// TaskInfo 方法
func (t *TaskInfo) addMessage(role, content string) {
	t.Messages = append(t.Messages, Message{
		Role:    role,
		Content: content,
	})
}

func (t *TaskInfo) getRecentMessages(limit int) []Message {
	if len(t.Messages) <= limit {
		return t.Messages
	}
	return t.Messages[len(t.Messages)-limit:]
}

func (t *TaskInfo) addToolCall(toolCall ToolCallInfo) {
	t.ToolCalls = append(t.ToolCalls, toolCall)
}

func (t *TaskInfo) updateContext(key string, value interface{}) {
	t.Context[key] = value
}

// 辅助函数
func generateSessionID() string {
	return fmt.Sprintf("ai_session_%d", time.Now().UnixNano())
}

func generateMessageID() string {
	return fmt.Sprintf("msg_%d", time.Now().UnixNano())
}

func generateTaskID() string {
	return fmt.Sprintf("task_%d", time.Now().UnixNano())
}
